import os.path
import pickle
import random
from collections import defaultdict

import pandas as pd
from datetime import datetime

from shared_variables import CONN
import numpy as np

from utils.utils import calculate_penalty, calculate_earnings, calculate_physical_penalty, calculate_cash_retention, \
    LOGGER


class LevelData:
    def __init__(self, project_id, colleague_number=8):
        self.project_id = project_id
        self.data = {}  # 存储关卡数据的字典，以表名为键，DataFrame为值
        self.me_character_id = -1
        self.load_data(colleague_number)
        self.talk_role = ''
        self.talks = defaultdict(list)

    def load_data(self, colleague_number):
        # 加载其他表的数据
        tables = ['talks', 'task', 'relationships', 'taskPeopleLinks', 'projectTaskLinks', 'goods']
        for table in tables:
            query = f"SELECT * FROM {table}"
            self.data[table] = pd.read_sql_query(query, CONN)

        query = f"SELECT * FROM people"
        self.data['people'] = pd.read_sql_query(query, CONN).sample(colleague_number).reset_index(drop=True)

        query = f"SELECT * FROM project WHERE id = {self.project_id}"
        self.data['project'] = pd.read_sql_query(query, CONN)

        query = f"SELECT * FROM projectTaskLinks WHERE project_id = {self.project_id}"
        self.data['projectTaskLinks'] = pd.read_sql_query(query, CONN)

    def update_one_talk(self, role, name, content):
        temp_talk = {"project_id": self.project_id, "role": role, "name": name, "speak_time": datetime.now(),
                     "content": content}

        self.talks[self.talk_role].append(temp_talk)
        print(self.talks)

    def save_data_to_database(self):
        # 将缓存中的数据写回到数据库中
        for table, df in self.data.items():
            # 示例：使用pandas的to_sql方法将DataFrame写入数据库
            df.to_sql(table, CONN, if_exists='replace', index=False)

        # 提交事务并关闭数据库连接
        CONN.commit()

    def set_me_character(self, character_id):
        self.me_character_id = character_id

    def load_characters(self):
        return self.data['people']

    def set_me_character_id(self, self_id):
        self.me_character_id = self_id

    def get_me_info(self) -> pd.DataFrame:
        return self.data['people'][self.data['people']['role'] == 'me']

    def get_people_role(self) -> pd.DataFrame:
        my_name = self.get_me_info()['name'].values[0]
        filtered_df = self.data['relationships'][self.data['relationships']['peopleAB'].str.contains(my_name)].copy()
        # 去除姓名和逗号，只保留关系描述
        filtered_df.loc[:, 'name'] = filtered_df['peopleAB'].apply(lambda x: x.replace(my_name, '').replace(',', ''))
        result = pd.merge(self.data['people'], filtered_df[['name', 'relationship']], how='outer', on='name')
        return result

    def get_boss(self) -> pd.Series:
        people = self.get_people_role()
        boss = people[people['role'] == 'boss']
        return boss

    def set_talk_role(self, name):
        self.talk_role = name

    def initialize_random_role(self, me_id):
        roles = ['boss', 'employee']
        boss_assigned = False

        # 生成随机的role数据
        while True:
            role_data = np.random.choice(roles, size=self.data['people'].shape[0])
            if role_data[me_id] != 'boss':
                break

        # 确保只有一个boss
        for i, role in enumerate(role_data):
            if role == 'boss' and not boss_assigned:
                role_data[i] = 'boss'
                boss_assigned = True
            else:
                role_data[i] = 'employee'

        # 将me的role设为'me'
        role_data[self.data['people'].index == me_id] = 'me'
        self.data['people']['role'] = role_data

    def text_fix_role(self, me_id, boss_id):
        roles = ['employee'] * self.data['people'].shape[0]
        roles[me_id] = 'me'
        roles[boss_id] = 'boss'
        self.data['people']['role'] = roles

    def calc_cash(self, df):
        project_duration = self.data['project'].loc[self.data['project']['id'] == self.project_id, 'duration'].values[0]
        # 计算每个人名下的task_duration之和
        grouped_df = df.groupby('name_people').agg(
            {'bonus_task': 'sum', 'cost': 'sum', 'cash': 'mean', 'salary': 'mean', 'bonus': "mean",
             'energy': 'mean', 'rent_expense': 'mean', 'living_expense': 'mean', 'social_expense': 'mean'})
        # 找到最大的项目时间
        project_real_duration = grouped_df['cost'].max()
        # 计算惩罚系数
        penalty = project_duration / project_real_duration
        penalty = penalty if penalty < 1 else 1

        over_time = project_real_duration - project_duration
        # 计算工资收益
        grouped_df['salary_earnings'] = grouped_df['salary'] * project_real_duration / 30
        grouped_df['salary_earnings'] = grouped_df['salary_earnings'].astype(int)
        # 计算任务收益
        grouped_df['task_earnings'] = grouped_df['bonus_task'] * penalty
        grouped_df['task_penalty'] = grouped_df['bonus_task'] - grouped_df['task_earnings']

        grouped_df['physical_penalty_alpha'] = grouped_df['energy'] - grouped_df['cost']
        grouped_df.loc[grouped_df['physical_penalty_alpha'] <= 0, 'physical_penalty'] = grouped_df[
            'physical_penalty_alpha']
        grouped_df['physical_penalty'] = grouped_df['physical_penalty'].fillna(0)
        grouped_df['physical_penalty'] *= 50  # 体力惩罚系数
        # 计算现金留存
        grouped_df['cash_retention'] = grouped_df['cash'] + grouped_df['salary_earnings'] + grouped_df[
            'task_earnings'] - \
                                       grouped_df['rent_expense'] - grouped_df['living_expense'] - grouped_df[
                                           'social_expense'] - grouped_df['physical_penalty']
        # 保留其他数据
        grouped_df['project_overdue'] = over_time if over_time > 0 else 0

        # 返回结果
        return grouped_df

    def temp_init_task_arrangement(self):
        temp = pd.merge(self.data['project'], self.data['projectTaskLinks'], left_on='id', right_on='project_id')
        project_task_df = pd.merge(temp, self.data['task'], left_on='task_id', right_on='id', suffixes=('', '_task'))
        relationships = []

        # 随机生成对应关系
        tasks = project_task_df['id_task'].unique()
        people = self.data['people']['id'].to_list()
        np.random.shuffle(people)

        for i, task_id in enumerate(tasks):
            if i < len(people):
                people_id = people[i]
            else:
                random_index = np.random.randint(len(people))
                people_id = people[random_index]

            relationships.append({
                'project_id': project_task_df.loc[project_task_df['id_task'] == task_id, 'project_id'].values[0],
                'task_id': task_id,
                'people_id': people_id
            })

        relationship_df = pd.DataFrame(relationships)
        self.data['taskPeopleLinks'] = relationship_df
        return relationship_df

    def get_task_people(self, taskPeopleLinksDf=pd.DataFrame) -> pd.DataFrame:
        if taskPeopleLinksDf.empty:
            taskPeopleLinksDf = self.data['taskPeopleLinks']

        temp = pd.merge(self.data['project'], self.data['projectTaskLinks'], left_on='id', right_on='project_id')
        project_task_df = pd.merge(temp, self.data['task'], left_on='task_id', right_on='id', suffixes=('', '_task'))
        temp = pd.merge(project_task_df, taskPeopleLinksDf, left_on='id_task', right_on='task_id',
                        suffixes=('', '_'))
        result = pd.merge(temp, self.data['people'], left_on='people_id', right_on='id', how='outer',
                          suffixes=('', '_people'))
        return result

    def reassign_task(self, task_name, new_people):
        try:
            LOGGER.info(f'任务分配之前 {self.data["taskPeopleLinks"]}')
            people_id = self.data['people'].loc[self.data['people']['name'] == new_people, 'id'].values[0]
            task_id = self.data['task'].loc[self.data['task']['name'] == task_name, 'id'].values[0]
            self.data['taskPeopleLinks'].loc[
                self.data['taskPeopleLinks']['task_id'] == task_id, 'people_id'] = people_id
            LOGGER.info(f'任务重新分配之后 { self.data["taskPeopleLinks"]}')
        except:
            LOGGER.info(f'{task_name} 和 {new_people} 有一个没有匹配上')

    def predict_task(self, task_name, new_people):
        people_id = self.data['people'].loc[self.data['people']['name'] == new_people, 'id'].values[0]
        task_id = self.data['task'].loc[self.data['task']['name'] == task_name, 'id'].values[0]
        temp = self.data['taskPeopleLinks'].copy()
        temp.loc[temp['task_id'] == task_id, 'people_id'] = people_id
        return temp

    def get_relationship(self, person1, person2):
        key = ','.join(sorted([person1, person2]))
        print(key)
        return self.data['relationships'].loc[self.data['relationships']['peopleAB'] == key, 'relationship'].values[0]

    def update_relationship(self, person1, person2, new_relationship):
        key = ','.join(sorted([person1, person2]))
        self.data['relationships'].loc[self.data['relationships']['peopleAB'] == key, 'relationship'] = new_relationship

    def predict_finance(self, data):
        df = self.get_task_people()
        cash_current = self.calc_cash(df)

        single_name = data['colleague_tasks'].split('和')[0]
        new_taskPeopleLinks = self.predict_task(single_name, data['my_name'])
        df = self.get_task_people(new_taskPeopleLinks)
        future_cash = self.calc_cash(df)
        current_financial_status = cash_current.loc[data['colleague_name'], 'cash_retention']
        future_financial_status = future_cash.loc[data['colleague_name'], 'cash_retention']
        return current_financial_status, future_financial_status

    def prepare_data(self, colleague_name=None):
        me = self.get_me_info()
        tasks = self.get_task_people()
        peoples = self.get_people_role()
        project = self.data['project']

        me_describe = me['describe'].values[0]
        my_name = me['name'].values[0]

        task_me = '和'.join(tasks.loc[tasks['role'] == "me", 'name_task'].tolist())
        project_name = project.loc[project['id'] == self.project_id, 'name'].values[0]

        if colleague_name:
            if isinstance(colleague_name, str):
                colleague = peoples[peoples['name'] == colleague_name]
            else:
                colleague = peoples[peoples['id'] == colleague_name]
                colleague_name = colleague.name.values[0]

            colleague_tasks = '和'.join(tasks.loc[tasks['name_people'] == colleague_name, 'name_task'].tolist())
            colleague_describe = colleague['describe'].values[0]
            relationship = self.get_relationship(my_name, colleague_name)
            data = {
                "relationship": relationship,
                "me_describe": me_describe,
                "colleague_describe": colleague_describe,
                "my_task": task_me,
                "boss_name": "",
                "colleague_tasks": colleague_tasks,
                "project_name": project_name,
                "colleague_name": colleague_name,
                "my_name": my_name
            }
        else:
            boss = peoples[peoples['role'] == 'boss']
            boss_name = boss['name'].values[0]
            boss_describe = boss['describe'].values[0]
            relationship = self.get_relationship(my_name, boss_name)
            data = {
                "relationship": relationship,
                "me_describe": me_describe,
                "boss_describe": boss_describe,
                "colleague_tasks": "",
                "project_name": project_name,
                "my_task": task_me,
                "boss_name": boss_name,
                "colleague_name": '',
                "my_name": my_name
            }

        return data

    def prepare_data_between_colleagues(self, colleague_me, colleague_name):
        me = self.data['people'][self.data['people']['name'] == colleague_me]
        tasks = self.get_task_people()
        peoples = self.get_people_role()
        project = self.data['project']

        me_describe = me['describe'].values[0]
        my_name = me['name'].values[0]

        task_me = '和'.join(tasks.loc[tasks['role'] == "me", 'name_task'].tolist())
        project_name = project.loc[project['id'] == self.project_id, 'name'].values[0]

        if isinstance(colleague_name, str):
            colleague = peoples[peoples['name'] == colleague_name]
        else:
            colleague = peoples[peoples['id'] == colleague_name]
            colleague_name = colleague.name.values[0]

        os.path.join(os.curdir, 'assets/head',  f'{my_name}.png')
        colleague_tasks = '和'.join(tasks.loc[tasks['name_people'] == colleague_name, 'name_task'].tolist())
        colleague_describe = colleague['describe'].values[0]
        relationship = self.get_relationship(my_name, colleague_name)
        data = {
            "relationship": relationship,
            "me_describe": me_describe,
            "colleague_describe": colleague_describe,
            "my_task": task_me,
            "avatar_me": os.path.join(os.curdir, 'assets/head',  f'{my_name}.png'),
            "avatar_colleague": os.path.join(os.curdir, 'assets/head',  f'{colleague_name}.png'),
            "boss_name": "",
            "colleague_tasks": colleague_tasks,
            "project_name": project_name,
            "colleague_name": colleague_name,
            "my_name": my_name
        }
        return data

if __name__ == '__main__':
    # 创建一个LevelData实例并加载数据
    level = LevelData(project_id=1)
    level.temp_init_task_arrangement()
    data = level.prepare_data()
    cc = level.predict_finance(data)
    a = 0
